import os

from configs import QIANFAN_AK, QIANFAN_SK
from langchain.chat_models import QianfanChatEndpoint
from langchain_core.messages import HumanMessage, SystemMessage

"""使用 Chat Models 来对话"""
# 01-first-llm.py LLM例子都是使用的text in, text out方式，直接向LLM输入prompt text，让其返回generation text.
# 此外，LangChain 还提供对话模型，它是LLM的一个变种，其底层还是LLM，只是其编程接口与text in, text out的风格有所不同，它使用的是chat messages作为输入和输出。

os.environ["QIANFAN_AK"] = QIANFAN_AK
os.environ["QIANFAN_SK"] = "3yGoXGIdCCy8QLwcLCLdh3UGGlM7PKNn"

chatllm = QianfanChatEndpoint()
# result = chatllm([HumanMessage(content='请把以下句子翻译成中文: suzhou is a greate city')])
# # AIMessage
# print(type(result))
#
# # content='苏州是一个美丽的城市。'
# # additional_kwargs={'id': 'as-it3v92s6bf', 'object': 'chat.completion', 'created': 1704104582, 'result': '苏州是一个美丽的城市。', 'is_truncated': False, 'need_clear_history': False, 'usage': {'prompt_tokens': 15, 'completion_tokens': 5, 'total_tokens': 20}}
# print(result)
# print(result.content)
#
# # 会有token的数据
# print(result.additional_kwargs)

message = [
    SystemMessage(content='你是一个将英语翻译成中文的助手'),
    HumanMessage(content='suzhou is a greate city')
]
response = chatllm(message)

# content='苏州是一个美丽的城市。'
# additional_kwargs={'id': 'as-gsp73cn4bb', 'object': 'chat.completion', 'created': 1704105870, 'result': '苏州是一个美丽的城市。', 'is_truncated': False, 'need_clear_history': False, 'usage': {'prompt_tokens': 18, 'completion_tokens': 5, 'total_tokens': 23}}
print(response)

batch_message = [
    [
        SystemMessage(content='你是一个将英语翻译成中文的助手'),
        HumanMessage(content='suzhou is a greate city')
    ],
    [
        SystemMessage(content='你是一个将中文翻译成英语的助手'),
        HumanMessage(content='苏州是一个美丽的城市')
    ]
]
#
# # 可以一次传多组信息
llm_result = chatllm.generate(batch_message)

# [ChatGeneration(text='苏州是一个美丽的城市，有许多著名的景点和文化遗产。', generation_info={'finish_reason': 'stop', 'id': 'as-zuetdrwx5b', 'object': 'chat.completion', 'created': 1704106037, 'result': '苏州是一个美丽的城市，有许多著名的景点和文化遗产。', 'is_truncated': False, 'need_clear_history': False, 'usage': {'prompt_tokens': 18, 'completion_tokens': 11, 'total_tokens': 29}}, message=AIMessage(content='苏州是一个美丽的城市，有许多著名的景点和文化遗产。', additional_kwargs={'id': 'as-zuetdrwx5b', 'object': 'chat.completion', 'created': 1704106037, 'result': '苏州是一个美丽的城市，有许多著名的景点和文化遗产。', 'is_truncated': False, 'need_clear_history': False, 'usage': {'prompt_tokens': 18, 'completion_tokens': 11, 'total_tokens': 29}}))]
# [ChatGeneration(text='苏州是一个美丽的城市。', generation_info={'finish_reason': 'stop', 'id': 'as-h6b7091rxa', 'object': 'chat.completion', 'created': 1704106037, 'result': '苏州是一个美丽的城市。', 'is_truncated': False, 'need_clear_history': False, 'usage': {'prompt_tokens': 13, 'completion_tokens': 5, 'total_tokens': 18}}, message=AIMessage(content='苏州是一个美丽的城市。', additional_kwargs={'id': 'as-h6b7091rxa', 'object': 'chat.completion', 'created': 1704106037, 'result': '苏州是一个美丽的城市。', 'is_truncated': False, 'need_clear_history': False, 'usage': {'prompt_tokens': 13, 'completion_tokens': 5, 'total_tokens': 18}}))]

for generation in llm_result.generations:
    print(generation)
#
# # 同样可以通过llm_output获取token使用情况：
# print(llm_result.llm_output)

"""BaseMessage"""
# 对话模型接受一个Message的数组作为输入,现在支持的Message类型包括：
# AIMessage: AI模型输出的Message
# HumanMessage: 用户输入的Message，往往作为in-context learning的输入，给AI做few-shot
# SystemMessage: 设定AI模型的role用到的Message
# ToolMessage: 将执行工具的结果传递回模型的消息。
# ChatMessage: 接受任意角色参数


"""BaseChatModel"""
# BaseChatModel 继承 BaseLanguageModel
# invoke -> generate_prompt -> generate -> _generate_with_cache -> _generate
# __call__ -> generate -> _generate_with_cache -> _generate
# predict -> __call__
# apredict -> _call_async -> agenerate -> _agenerate_with_cache -> _agenerate
# stream -> _stream

"""SimpleChatModel"""
# SimpleChatModel 继承 BaseChatModel
# 实现了 _generate 方法 -> _call
# 定义了 _call 方法

"""QianfanChatEndpoint"""
# QianfanChatEndpoint 继承 BaseChatModel
# 实现了 _generate 方法 -> _stream
# 实现了 _agenerate 方法 -> _astream
# 实现了 _stream 方法
# 实现了 _astream 方法
